package com.xtyhz.service.impl;

import cn.hutool.json.JSONUtil;
import com.xtyhz.service.DeepSeekR1;
import jakarta.servlet.ServletOutputStream;
import jakarta.servlet.http.HttpServletResponse;
import okhttp3.*;
import org.springframework.stereotype.Service;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.net.SocketTimeoutException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

@Service
public class DeepSeekR1Impl implements DeepSeekR1 {

    @Override
    public String deepSeek(String prompt, HttpServletResponse res) {
        Map<String, Object> params = new HashMap<>();
        params.put("model", "deepseek-r1:7b");

        Map<String, String> message = new HashMap<>();
        message.put("role", "user");
        message.put("content", prompt);

        List<Map<String,String>> messages = new ArrayList<>();
        messages.add(message);
        params.put("messages", messages);
        params.put("stream", true);

        String jsonParams = JSONUtil.toJsonStr(params);

        // 增加超时设置（关键修改）
        OkHttpClient client = new OkHttpClient().newBuilder()
                .connectTimeout(30, TimeUnit.SECONDS)  // 连接超时
                .readTimeout(100, TimeUnit.SECONDS)     // 读取超时（流式响应需要更长时间）
                .writeTimeout(60, TimeUnit.SECONDS)    // 写入超时
                .build();

        MediaType mediaType = MediaType.parse("application/json");
        RequestBody body = RequestBody.create(jsonParams, mediaType);

        Request request = new Request.Builder()
                .url("http://localhost:11434/v1/chat/completions")
                .post(body)
                .addHeader("Content-Type", "application/json")
                .addHeader("Accept", "application/json")
//                .addHeader("Authorization", "Bearer sk-22f0a8a4af7343d8b13aa84630a180ee")
                .build();

        // 使用try-with-resources确保资源释放
        try (Response response = client.newCall(request).execute();
             ServletOutputStream outputStream = res.getOutputStream()) {

            res.setContentType("text/event-stream");
            res.setCharacterEncoding("UTF-8");
            res.setHeader("Cache-Control", "no-cache");

            if (!response.isSuccessful()) {
                // 处理错误响应
                res.setStatus(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
                outputStream.write(("HTTP error: " + response.code()).getBytes());
                return "error";
            }

            ResponseBody responseBody = response.body();
            if (responseBody != null) {
                try (InputStream inputStream = responseBody.byteStream()) {
                    BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream));
                    String line;
                    // 使用缓冲读取方式逐行处理
                    while ((line = reader.readLine()) != null) {
                        if (line.startsWith("data: ")) {
                            String eventData = line.substring(6).trim();
                            if (!eventData.isEmpty()) {
                                outputStream.write(("data: " + eventData + "\n\n").getBytes());
                                outputStream.flush();
                            }
                        }
                    }
                }
            }
        } catch (SocketTimeoutException ste) {
            // 处理超时异常
            res.setStatus(HttpServletResponse.SC_GATEWAY_TIMEOUT);
            return "timeout_error";
        } catch (IOException e) {
            // 处理其他IO异常
            res.setStatus(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
            return "io_error";
        }

        return "success";
    }
}
